/*
 * Copyright 2008 Neil Bacon
 * Distributed under the terms of the Mozilla Public License version 1.1
 * or the GNU General Public License version 3.
 */

package nb.xadd.operator.impl;

import static java.lang.Double.isNaN;

import java.util.BitSet;

import nb.cache.CacheMap;
import nb.xadd.XADDRuntimeException;
import nb.xadd.node.Node;
import nb.xadd.node.NodeFactory;
import nb.xadd.node.impl.BooleanDecisionNode;
import nb.xadd.node.impl.ConstantValueNode;
import nb.xadd.node.impl.RealEqualDecisionNode;
import nb.xadd.node.impl.RealGreaterOrEqualDecisionNode;
import nb.xadd.operator.BinaryOperator;
import nb.xadd.operator.OperatorFactory;
import nb.xadd.operator.UnaryOperator;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

/**
 * Integration over the variables specified by realVarIndexSet.
 * All terminal nodes must be ConstantValueNodes or a MDPRuntimeException is thrown.
 * TODO: handle LinearValueNodes and PolynomialValueNodes.
 */
public class Integration implements UnaryOperator {

    private Log log = LogFactory.getLog(getClass());

    private final NodeFactory nodeFactory;
    // private final OperatorFactory operatorFactory;
    private final CacheMap<UnaryOperatorCacheKey, Node> unaryOperatorCache;
    private final BitSet booleanVarIndexSet;
    private final BitSet realVarIndexSet;
    private final int hashCode;
    private final BinaryOperator plus;
    private final BinaryOperator times;

    protected Integration(NodeFactory nodeFactory, OperatorFactory operatorFactory, CacheMap<UnaryOperatorCacheKey, Node> unaryOperatorCache, BitSet booleanVarIndexSet, BitSet realVarIndexSet) {
        super();
        this.nodeFactory = nodeFactory;
        // this.operatorFactory = operatorFactory;
        this.unaryOperatorCache = unaryOperatorCache;
        this.booleanVarIndexSet = booleanVarIndexSet;
        this.realVarIndexSet = realVarIndexSet;
        hashCode = 29 * booleanVarIndexSet.hashCode() + realVarIndexSet.hashCode();
        plus = operatorFactory.getAddition();
        times = operatorFactory.getMultiplication();
    }

    public Node execute(Node n) {
        UnaryOperatorCacheKey key = new UnaryOperatorCacheKey(this, n);
        Node result = unaryOperatorCache.get(key);
        if (result == null) {
            result = walk(n, -1, Double.NaN, Double.NaN);
            unaryOperatorCache.put(key, result);
        }
        return result;
    }

    private Node walk(Node n, int realVarIndex, double min, double max) {
        Node result = null;
        if (n instanceof RealGreaterOrEqualDecisionNode) {
            RealGreaterOrEqualDecisionNode dn = (RealGreaterOrEqualDecisionNode) n;
            if (dn.getRealVarIndex() == realVarIndex) {
                if (log.isDebugEnabled()) {
                    log.debug("walk: true range " + dn.getTestValue() + " <= x" + dn.getRealVarIndex() + " < " + max);
                    log.debug("walk: false range " + min + " <= x" + dn.getRealVarIndex() + " < " + dn.getTestValue());
                }
                Node trueChild = walk(dn.getTrueChild(), dn.getRealVarIndex(), dn.getTestValue(), max);
                Node falseChild = walk(dn.getFalseChild(), dn.getRealVarIndex(), min, dn.getTestValue());
                return trueChild.apply(plus, falseChild);
            } else if (realVarIndexSet.get(dn.getRealVarIndex())) {
                if (log.isDebugEnabled()) {
                    log.debug("walk: integrate x" + dn.getRealVarIndex() + " split at " + dn.getTestValue());
                }
                Node trueChild = walk(dn.getTrueChild(), dn.getRealVarIndex(), dn.getTestValue(), Double.NaN);
                Node falseChild = walk(dn.getFalseChild(), dn.getRealVarIndex(), Double.NaN, dn.getTestValue());
                result = trueChild.apply(plus, falseChild);
            } else {
                if (log.isDebugEnabled()) {
                    log.debug("walk: preserve node x" + dn.getRealVarIndex() + " >= " + dn.getTestValue() + " ?");
                }
                // use execute rather than walk(node, -1, Double.NaN, Double.NaN) so that intermediate results are cached
                Node trueChild = execute(dn.getTrueChild());
                Node falseChild = execute(dn.getFalseChild());
                result = nodeFactory.getRealGreaterOrEqualDecisionNode(dn.getRealVarIndex(), dn.getTestValue(), trueChild, falseChild);
            }
        } else if (n instanceof RealEqualDecisionNode) {
            RealEqualDecisionNode dn = (RealEqualDecisionNode) n;
            // use execute rather than walk(node, -1, Double.NaN, Double.NaN) so that intermediate results are cached
            Node trueChild = execute(dn.getTrueChild());
            Node falseChild = execute(dn.getFalseChild());
            if (realVarIndexSet.get(dn.getRealVarIndex())) {
                if (log.isDebugEnabled()) {
                    log.debug("walk: sum result for each child of node x" + dn.getRealVarIndex() + " == " + dn.getTestValue() + " ?");
                }
                result = trueChild.apply(plus, falseChild);
            } else {
                if (log.isDebugEnabled()) {
                    log.debug("walk: preserve node x" + dn.getRealVarIndex() + " >= " + dn.getTestValue() + " ?");
                }
                result = nodeFactory.getRealGreaterOrEqualDecisionNode(dn.getRealVarIndex(), dn.getTestValue(), trueChild, falseChild);
            }
        } else if (n instanceof BooleanDecisionNode) {
            BooleanDecisionNode dn = (BooleanDecisionNode) n;
            // use execute rather than walk(node, -1, Double.NaN, Double.NaN) so that intermediate results are cached
            Node trueChild = execute(dn.getTrueChild());
            Node falseChild = execute(dn.getFalseChild());
            if (booleanVarIndexSet.get(dn.getBooleanVarIndex())) {
                if (log.isDebugEnabled()) {
                    log.debug("walk: sum result for each child of node b" + dn.getBooleanVarIndex() + " ?");
                }
                result = trueChild.apply(plus, falseChild);
            } else {
                if (log.isDebugEnabled()) {
                    log.debug("walk: preserve node b" + dn.getBooleanVarIndex() + " ?");
                }
                result = nodeFactory.getBooleanDecisionNode(dn.getBooleanVarIndex(), trueChild, falseChild);
            }
        } else if (n instanceof ConstantValueNode) {
            result = n;
        } else {
            throw new XADDRuntimeException("Integration can't handle Node type " + n.getClass().getName());
        }
        if (!isNaN(min) && !isNaN(max)) {
            result = result.apply(times, nodeFactory.getConstantValueNode(max - min));
        }
        return result;
    }

    @Override
    public int hashCode() {
        return hashCode;
    }

    @Override
    public boolean equals(Object obj) {
        if (this == obj)
            return true;
        if (obj == null || getClass() != obj.getClass())
            return false;
        Integration other = (Integration) obj;
        return realVarIndexSet.equals(other.realVarIndexSet);
    }

}
